5.2 행렬 믹서(Matrix Mixer) 이론과 블록 분해 알고리즘

5.2 행렬 믹서(Matrix Mixer) 이론과 블록 분해 알고리즘

본 장에서는 Mamba-2 아키텍처의 핵심 이론적 기반이 되는 행렬 믹서(Matrix Mixer) 이론과, 이를 최신 하드웨어 가속기 상에서 효율적으로 구현하기 위한 블록 분해(Block Decomposition) 알고리즘에 대해 심도 있게 다룬다. Mamba-1이 선택적 스캔(Selective Scan)을 통해 순차적 모델링의 효율성을 입증하며 포스트 트랜스포머 시대의 서막을 알렸다면, Mamba-2는 구조적 상태 공간 쌍대성(SSD: Structured State Space Duality) 프레임워크를 통해 상태 공간 모델(SSM)을 행렬 연산의 관점에서 완전히 재해석하였다. 이러한 재해석은 단순히 이론적 우아함을 넘어서, 최신 GPU 가속기(Tensor Core)의 연산 능력을 극대화하는 동시에 선형 시간 복잡도(O(T))를 유지하는 새로운 알고리즘의 설계를 가능케 했다. 본 절에서는 반분리 행렬(Semiseparable Matrix)의 수학적 성질부터 시작하여, 이를 블록 단위로 분해하여 병렬 처리를 수행하는 SSD 알고리즘의 작동 원리를 수식과 시스템적 관점에서 포괄적으로 분석한다.

1. 행렬 믹서(Matrix Mixer) 관점에서의 시퀀스 모델링

시퀀스 모델링의 본질은 입력 시퀀스 X \in \mathbb{R}^{T \times P}를 출력 시퀀스 Y \in \mathbb{R}^{T \times P}로 변환하는 함수를 학습하는 과정으로 정의할 수 있다. 심층 신경망의 발전 역사에서 이러한 변환은 순환 신경망(RNN), 합성곱 신경망(CNN), 그리고 트랜스포머(Transformer) 등 다양한 아키텍처 형태로 구현되어 왔으나, SSD 프레임워크는 이들을 행렬 믹서(Matrix Mixer) 또는 행렬 시퀀스 변환(Matrix Sequence Transformation) 이라는 하나의 통합된 관점에서 바라본다.1

1.1 행렬 변환(Matrix Transformation)의 보편적 정의

가장 일반적인 형태의 선형 시퀀스 변환은 입력 벡터 x에 대해 선형 연산자 M을 적용하여 출력 y를 얻는 선형 대수적 연산으로 표현된다.
y = Mx
여기서 x와 y는 시퀀스 길이 T와 채널 차원 P를 가지는 텐서를 평탄화(flatten)하거나 특정 차원에 대해 연산하는 벡터로 간주되며, M \in \mathbb{R}^{T \times T}는 시퀀스 믹싱(Sequence Mixing) 또는 토큰 믹싱(Token Mixing)을 수행하는 변환 행렬이다. 중요한 점은 이 M 행렬의 구조(Structure) 가 곧 모델의 귀납적 편향(Inductive Bias)과 연산 특성을 결정짓는다는 사실이다.2

이러한 관점에서 기존의 주요 시퀀스 모델들을 재분류하면 다음과 같은 구조적 특징이 드러난다.

  • 합성곱(Convolution): M이 토플리츠(Toeplitz) 행렬 구조를 가질 때, 이는 시불변(Time-invariant) 필터를 적용하는 합성곱 연산과 등가이다. 토플리츠 행렬은 대각선 성분이 일정하므로 파라미터 공유가 일어나며, 이는 국소적 특징 추출에 유리하다.
  • 어텐션(Attention): M이 입력 x에 의존하여 동적으로 생성되는 조밀한(Dense) 행렬일 때, 이는 자기 주의(Self-Attention) 메커니즘이 된다. 표준적인 어텐션에서 M = \text{softmax}(QK^\top)로 정의되며, 이는 모든 토큰 간의 쌍방향 상호작용을 모델링할 수 있게 한다. 그러나 M을 명시적으로 구체화(Materialize)할 경우 T \times T 크기의 메모리와 O(T^2)의 연산 비용이 발생한다.2
  • SSM(State Space Model): M이 반분리(Semiseparable) 구조를 가질 때, 이는 상태 공간 모델의 재귀적(Recurrent) 연산과 수학적으로 등가이다. 반분리 행렬은 대각선 아래의 부분 행렬들이 낮은 랭크(Low-Rank)를 가진다는 특징이 있으며, 이는 긴 시퀀스를 효율적으로 압축하여 전달하는 SSM의 특성을 행렬 구조로 표현한 것이다.5

Mamba-2에서 제안하는 SSD 계층은 이 M 행렬을 구성함에 있어, SSM의 재귀적 효율성과 어텐션의 병렬 처리 효율성을 동시에 만족시키는 최적의 구조를 탐색한다. 즉, M을 구조화된 반분리 행렬(Structured Semiseparable Matrix) 로 정의함으로써, 순차적 연산(Recurrence)과 병렬 행렬 연산(Matrix Multiplication)이라는 두 가지 상반된 계산 경로(Dual Computation Paths)를 수학적으로 연결한다.

1.2 반분리 행렬(Semiseparable Matrix)의 수학적 성질

SSD 알고리즘의 핵심을 이해하기 위해서는 반분리 행렬의 엄밀한 정의와 그 성질을 파악해야 한다. 반분리 행렬은 행렬의 구조적 랭크(Structured Rank) 특성에 의해 정의된다.3

정의 3.1 (N-반분리 행렬):

하삼각 행렬(Lower Triangular Matrix) M \in \mathbb{R}^{T \times T}이 N-반분리 행렬(N-Semiseparable Matrix) 이라는 것은, 대각선 포함 그 아래(lower triangular portion)에 존재하는 모든 부분 행렬(submatrix)의 랭크(Rank)가 최대 N임을 의미한다. 여기서 N은 상태 공간 모델의 상태 차원(State Dimension)에 해당한다.

이러한 랭크 제약은 행렬 M이 매우 효율적인 파라미터화(Parameterization)를 가질 수 있음을 시사한다. 구체적으로, 임의의 N-반분리 행렬은 순차적 반분리(Sequentially Semiseparable, SSS) 표현형을 가지며, 그 성분 M_{ij} (i \ge j)는 다음과 같이 분해된다.3
M_{ij} = C_i^\top A_{i:j}^{\times} B_j
이 수식의 구성 요소를 상태 공간 모델의 관점에서 해석하면 다음과 같다.

  • C_i \in \mathbb{R}^N: 시간 i에서의 출력 투영(Projection) 벡터.
  • B_j \in \mathbb{R}^N: 시간 j에서의 입력 투영 벡터.
  • A_{i:j}^{\times}: 시간 j에서 i까지의 상태 전이 행렬들의 곱(Product). 즉, A_{i:j}^{\times} = A_{i-1} A_{i-2} \cdots A_j이다.
  • N: 반분리 랭크(Order or Rank)이자 SSM의 은닉 상태(Hidden State) 차원이다.

이 정의는 상태 공간 모델의 시간 영역(Time-domain) 해와 정확히 일치한다. 이산 시간 SSM의 재귀식 h_t = A_t h_{t-1} + B_t x_t, y_t = C_t^\top h_t를 전개하여 출력 y를 입력 x에 대한 직접적인 함수로 표현하면, 그 변환 행렬은 필연적으로 위와 같은 반분리 구조를 갖게 된다. 이는 모든 선형 시불변(LTI) 또는 시변(LTV) SSM이 반분리 행렬 연산으로 해석될 수 있음을 증명한다.2

따라서, Mamba-2가 해결하고자 하는 문제는 “어떻게 하면 이 거대한 반분리 행렬 M을 명시적으로 생성하지 않으면서, 혹은 생성하더라도 매우 효율적으로 곱셈 연산을 수행할 것인가?“로 귀결된다.

2. 구조적 상태 공간 쌍대성(SSD)과 1-반분리 구조

Mamba-1을 비롯한 기존의 구조적 SSM(S4, H3 등)은 연산 효율성을 위해 A 행렬에 대각(Diagonal) 구조나 대각+저랭크(Diagonal Plus Low-Rank) 구조와 같은 제약을 가했다. Mamba-2는 여기서 한 걸음 더 나아가, 텐서 코어 활용과 어텐션 구조와의 연결을 위해 스칼라 구조(Scalar Structure) 라는 더 강력한 제약을 도입한다.4

2.1 스칼라 단위 행렬 제약 (Scalar Structure Constraint)

Mamba-2의 SSD 계층에서는 시간 t에서의 상태 전이 행렬 A_t를 스칼라 값과 단위 행렬의 곱으로 제한한다.
A_t = a_t I
여기서 a_t \in \mathbb{R}은 시간 t에 따라 변하는 스칼라 값이며, I는 N \times N 단위 행렬이다. 이 제약은 모델의 표현력을 일부 제한하는 것처럼 보일 수 있으나, 실제로는 다음과 같은 강력한 이점을 제공한다.

  1. 파라미터 공유 (Parameter Sharing): 상태 공간의 N개 요소가 모두 동일한 감쇠(Decay) 역학 a_t를 공유하게 된다. 이는 각 헤드(Head) 내의 모든 채널이 동일한 동역학을 따름을 의미하며, 이를 통해 메모리 대역폭 사용을 획기적으로 줄일 수 있다.4
  2. 가환성(Commutativity)과 행렬 분해: 스칼라 행렬은 모든 행렬과 교환 법칙이 성립한다. 이 성질 덕분에 복잡한 행렬 곱 형태의 전이 행렬 A_{i:j}^{\times}가 단순한 스칼라 곱으로 환원된다.

스칼라 구조 제약을 적용하면, 앞서 정의한 반분리 행렬의 성분 M_{ij}는 다음과 같이 극적으로 단순화된다.
M_{ij} = C_i^\top (a_{i-1} \cdots a_j I) B_j = (C_i^\top B_j) \cdot (a_{i-1} \cdots a_j)
이 식은 두 개의 독립적인 항의 요소별 곱(Hadamard Product, \circ)으로 완벽하게 분리된다.

  • 내적 항 (Dot Product Term): C_i^\top B_j. 이는 쿼리(Query)와 키(Key)의 내적을 계산하는 어텐션 메커니즘의 Q K^\top와 구조적으로 동일하다. 여기서 C는 쿼리, B는 키의 역할을 수행한다고 볼 수 있다.
  • 감쇠 항 (Decay/Mask Term): L_{ij} = \prod_{k=j}^{i-1} a_k. 이는 입력 j가 출력 i에 미치는 영향이 시간에 따라 어떻게 감쇠하는지를 나타내는 항이다. 이는 어텐션 메커니즘에서의 위치 인코딩(Positional Encoding) 또는 마스크(Mask) 행렬의 역할을 수행한다.

결과적으로, Mamba-2의 시스템 행렬 M은 다음과 같은 1-반분리 행렬(1-Semiseparable Matrix) 의 형태로 표현된다.
M = L \circ (C B^\top)
여기서 L은 L_{ij}를 성분으로 가지는 1-반분리 마스크 행렬이고, C B^\top는 랭크가 N인 저랭크 행렬이다. 이 수식은 SSM과 어텐션 사이의 수학적 쌍대성(Duality)을 명확하게 보여준다. 즉, 스칼라 구조를 가진 SSM은 요소별 마스크가 적용된 선형 어텐션(Linear Attention)과 수학적으로 정확히 등가이다.2

2.2 쌍대성(Duality)이 제공하는 두 가지 계산 경로

이러한 쌍대성은 모델을 상황에 따라 두 가지 다른 모드로 실행할 수 있는 유연성을 제공한다.

  1. 선형 모드 (Linear/Recurrent Mode):
  • 형태: h_t = a_t h_{t-1} + b_t x_t, y_t = c_t^\top h_t
  • 특징: 전통적인 SSM 또는 RNN의 실행 방식이다. 시간 복잡도는 O(T)이며, 이전 상태 h_{t-1}만 메모리에 유지하면 되므로 추론(Inference) 시 KV 캐시(KV Cache)가 필요 없고 메모리 사용량이 일정하다(Constant Memory). 이는 긴 시퀀스를 생성할 때 압도적인 효율성을 제공한다.3
  1. 이차 모드 (Quadratic/Attention Mode):
  • 형태: Y = (L \circ (C B^\top)) X
  • 특징: 전체 시퀀스를 한 번에 처리하는 방식이다. L과 C B^\top 행렬을 구체화하여 곱셈을 수행한다. 나이브하게 구현할 경우 O(T^2)의 연산량이 발생하지만, 최신 GPU의 텐서 코어는 순차적 연산보다 행렬 곱(Matmul)을 훨씬 빠르게 처리한다. 따라서 학습(Training) 시에는 병렬화가 가능한 이 모드가 유리할 수 있다.1

하지만, 단순한 이차 모드는 시퀀스 길이 T가 길어질수록 메모리와 연산 비용이 급증한다는 트랜스포머의 고질적인 문제를 그대로 답습하게 된다. 여기서 Mamba-2는 선형 모드의 효율성과 이차 모드의 하드웨어 친화성을 결합하기 위해 블록 분해(Block Decomposition) 라는 제3의 길을 제시한다.

3. 블록 분해 알고리즘 (Block Decomposition Algorithm)

순수 SSM(재귀) 방식은 GPU 병렬화가 어렵고, 순수 어텐션(이차 모드) 방식은 시퀀스 길이에 대해 이차적인 비용이 발생한다. 블록 분해 알고리즘은 전체 시퀀스를 고정된 크기의 청크(Chunk) 또는 블록으로 분할하여, 블록 내부(Intra-chunk)는 고속 행렬 연산으로 처리하고, 블록 간(Inter-chunk) 연결은 재귀적으로 처리하는 하이브리드 전략이다. 이는 하드웨어 효율성과 알고리즘 복잡도 사이의 최적점을 찾아낸 결과이다.7

3.1 행렬 M의 블록 구조화 및 인수분해

전체 변환 행렬 M을 크기 Q \times Q의 블록들로 격자처럼 분할한다고 가정하자. M은 하삼각 행렬이므로, 이 블록들은 위치에 따라 크게 세 가지 유형으로 분류되며, 각각 다른 전략으로 계산된다.7

표 5.2.1 반분리 행렬의 블록 유형 및 계산 전략

블록 유형위치구조적 특징계산 전략 (알고리즘 단계)
대각 블록 (Diagonal Blocks)주대각선 상 (i=j)더 작은 크기(Q \times Q)의 1-반분리 행렬 (Causal)2차 모드 (Attention): 로컬 어텐션 연산 수행 (Step 1)
비대각 블록 (Off-diagonal)대각선 하단 (i > j)완전한 저랭크 (Low-Rank) 행렬인수분해 (Factorization): 상태 전달을 통한 계산 (Step 2, 4)
상단 블록 (Upper Triangular)대각선 상단 (i < j)영행렬 (Zero Matrix)계산 불필요 (Causality)

반분리 행렬의 정의에 의해, 대각선 아래에 위치한 비대각 블록들은 랭크 N을 가지는 저랭크 행렬로 인수분해될 수 있다. 즉, i번째 행 블록과 j번째 열 블록 (i > j) 사이의 상호작용을 나타내는 부분 행렬 M_{block}^{(i,j)}는 다음과 같이 분해된다.
M_{block}^{(i,j)} = L_{block}^{(i,j)} \cdot (C_{block}^{(i)} B_{block}^{(j)\top})
이러한 구조적 특성을 활용하여 SSD 알고리즘은 전체 T \times T 행렬을 구체화하지 않고도, 블록 단위의 연산을 통해 정확한 결과를 산출한다.

3.2 SSD 알고리즘의 4단계 실행 과정 (The 4 Steps)

SSD 알고리즘은 입력 X를 길이 Q의 청크들로 분할한 후, 다음 4단계의 과정을 거쳐 출력을 계산한다. 이 과정은 텐서 코어(Tensor Core)를 활용하기 위해 행렬 곱(Matmul)을 적극적으로 사용하면서도, 전체 복잡도를 시퀀스 길이에 대해 선형(O(T))으로 유지한다.7

3.2.1 단계: 청크 내부 출력 계산 (Intra-chunk Outputs)

첫 번째 단계는 각 청크 내에서의 로컬 출력을 계산하는 것이다. 이는 “이전 청크로부터 넘어온 상태가 0이라고 가정했을 때(Initial state = 0), 현재 청크의 입력이 만드는 출력“을 구하는 것과 같다.

  • 수학적 해석: 각 대각 블록(Diagonal Block)에 대해 이차 모드(Attention-like) 연산을 수행한다. 대각 블록 M_{diag}는 Q \times Q 크기의 작은 1-반분리 행렬이므로, 이를 명시적으로 계산해도 비용이 크지 않다.
  • 연산: Y_{\text{diag}} = (L_{\text{local}} \circ (C_{\text{local}} B_{\text{local}}^\top)) X_{\text{local}}
  • 구현: 코드 상에서는 torch.einsum을 사용하여 배치 처리된 행렬 곱으로 구현된다. L_{\text{local}}은 마스크 행렬이며, Mamba-2에서는 이를 명시적으로 Q \times Q 행렬로 만들지 않고 segsum 연산을 통해 효율적으로 계산한다.
  • 하드웨어: 각 청크가 독립적이므로 완벽한 병렬 처리가 가능하며, 텐서 코어의 고밀도 연산 능력을 100% 활용한다.

3.2.2 단계: 청크 상태 계산 (Chunk States)

두 번째 단계는 각 청크의 끝에서 생성되는 최종 상태(Final State)를 계산하는 것이다. 이는 “초기 상태가 0일 때, 해당 청크의 입력들이 축적되어 만들어낸 최종 상태“를 의미한다.

  • 수학적 해석: 반분리 행렬의 비대각(Off-diagonal) 블록이 저랭크임을 이용한다. 비대각 블록은 C 항과 B 항의 곱으로 분해되는데, 여기서 입력 측인 B 항과 감쇠(Decay) 항을 집계하여 상태 h를 형성한다.
  • 연산: h_{\text{chunk}} = \sum_{t=1}^{Q} (A_{t:Q}^{\times} B_t x_t)
  • 구현: 이 단계 역시 torch.einsum을 통한 행렬 곱으로 처리된다. 입력 X와 입력 투영 B를 결합하고, 청크 내에서의 누적 감쇠를 적용하여 압축된 상태 정보를 생성한다.

3.2.3 단계: 청크 간 상태 전달 (Pass States / Inter-chunk Recurrence)

세 번째 단계는 이전 단계에서 계산된 각 청크의 “로컬 최종 상태“들을 연결하여, 전체 시퀀스를 관통하는 “글로벌 상태“를 계산하는 과정이다. 이 단계가 바로 블록 간의 정보 흐름을 담당한다.

  • 수학적 해석: 청크 단위의 SSM 재귀(Recurrence)를 수행한다. 즉, 각 청크를 하나의 거대한 타임스텝으로 간주하고, 청크 간의 상태 전이를 계산한다.

  • 연산: h_{\text{global}}^{(k)} = A_{\text{chunk}}^{(k)} h_{\text{global}}^{(k-1)} + h_{\text{chunk}}^{(k)}

여기서 A_{\text{chunk}}^{(k)}는 k번째 청크 전체를 통과했을 때의 총 감쇠율(Total Decay)을 나타낸다.

  • 구현: 이 단계는 알고리즘 중 유일하게 순차적(Sequential)인 스캔 작업이 필요하다. 하지만, 전체 시퀀스 길이 T에 대해 수행하는 것이 아니라, 청크의 개수 T/Q에 대해서만 수행한다. Q가 일반적으로 64~128 정도로 설정되므로, 스캔해야 할 길이는 원래 길이의 1/64 ~ 1/128 수준으로 대폭 줄어든다. 따라서 이 단계가 전체 연산 시간에서 차지하는 비중은 극히 미미하며, 순차 처리로 인한 병목 현상이 거의 발생하지 않는다.7

3.2.4 단계: 출력 상태 보정 (Output States)

마지막 단계는 3단계에서 계산된 올바른 초기 상태(Global State)를 각 청크에 반영하여 최종 출력을 보정하는 것이다.

  • 수학적 해석: 각 청크의 출력 투영 파라미터 C와, 이전 청크들로부터 전달받은 글로벌 상태를 곱하여, 1단계에서 구한 로컬 출력에 더해준다. 이는 비대각 블록의 “좌측 항” 계산에 해당한다.
  • 연산: Y_{\text{final}} = Y_{\text{diag}} + Y_{\text{off-diag}} = Y_{\text{diag}} + C_{\text{local}} \cdot (A_{\text{decay}} h_{\text{prev\_global}})
  • 구현: 이 연산 또한 대규모 행렬 곱(Matmul)으로 처리되며, 각 청크별로 독립적인 계산이 가능하므로 완전한 병렬 실행이 가능하다.

3.3 블록 분해의 시스템적 이점

이 4단계 알고리즘은 겉보기에는 복잡해 보일 수 있으나, 실제 파이토치(PyTorch) 코드로 구현하면 약 25줄 내외의 매우 간결한 형태가 된다.8 이 알고리즘의 핵심 가치는 “텐서 코어가 좋아하는 연산(Matmul)은 최대로 늘리고, 싫어하는 연산(Scan/Recurrence)은 최소로 줄이는 것” 에 있다.

  • Matmul 비율의 극대화: 알고리즘의 1, 2, 4단계는 모두 고밀도 행렬 곱셈이다. 이는 GPU의 텐서 코어를 지속적으로 가동시켜(Compute Bound), 메모리 대역폭에 의해 성능이 제한되는(Memory Bound) 기존 스캔 알고리즘의 한계를 극복한다.
  • 메모리 IO 감소: 블록 내부 연산을 융합(Fusion)하고, 상태 전달을 위한 최소한의 정보만 메모리에 기록함으로써 HBM(High Bandwidth Memory)과 SRAM 사이의 데이터 이동을 최소화한다.

4. 계산 복잡도 분석과 하드웨어 효율성

SSD 알고리즘의 효율성은 이론적 복잡도 수치뿐만 아니라, 실제 하드웨어 아키텍처(GPU)의 특성과 얼마나 부합하느냐에 달려 있다.

4.1 시간 복잡도 (Time Complexity)

전체 시퀀스 길이를 T, 청크 크기를 Q, 상태 차원(State Dimension)을 N, 헤드 차원(Head Dimension)을 P라 할 때, 각 단계의 복잡도는 다음과 같다.

  1. Matmul 파트 (Step 1, 2, 4): O(T \cdot Q \cdot P + T \cdot N \cdot P)의 연산량을 가진다. 기존 어텐션이 O(T^2 P)인 것에 비하면, 시퀀스 길이 T에 대해 선형적이다. Q는 상수(예: 64)로 고정되므로 선형 복잡도가 유지된다.
  2. Recurrence 파트 (Step 3): O((T/Q) \cdot N^2) 또는 A가 대각 구조일 경우 O((T/Q) \cdot N)이다. 이는 T에 대해 선형적이면서도 그 계수가 1/Q로 매우 작아, 전체 연산량에서 무시할 수 있는 수준이다.

4.2 텐서 코어 활용 (Tensor Core Utilization)

NVIDIA A100 GPU를 기준으로, 텐서 코어를 사용한 BF16 행렬 연산 성능(312 TFLOPS)은 일반 FP32 산술 연산 성능(19 TFLOPS)보다 약 16배 이상 빠르다. H100에서는 그 격차가 더 벌어진다.8

  • Mamba-1의 한계: Mamba-1의 선택적 스캔은 요소별(Element-wise) 연산과 누적 합 위주였기 때문에, 텐서 코어를 제대로 활용하지 못하고 CUDA 코어의 성능에 의존했다.
  • Mamba-2의 도약: Mamba-2의 SSD 알고리즘은 연산의 대부분을 Q \times Q 또는 Q \times N 크기의 행렬 곱으로 변환한다. 이는 텐서 코어의 처리량을 최대로 끌어올려, Mamba-1 대비 학습 속도를 2배에서 최대 8배까지 향상시킨다.5

4.3 상태 차원 확장 (State Dimension Scaling)

효율적인 블록 분해 덕분에 Mamba-2는 상태 차원 N을 크게 늘릴 수 있다. Mamba-1은 연산 비용 문제로 보통 N=16을 사용했으나, Mamba-2는 N=64, 128 심지어 그 이상으로 확장 가능하다.9 상태 차원의 확장은 모델의 “기억 용량(Memory Capacity)“을 직접적으로 증대시키며, 이는 ’연상 회상(Associative Recall)’이나 ’다중 쿼리 추적(Multi-query Tracking)’과 같이 긴 문맥 속에서 세밀한 정보를 기억해야 하는 작업에서의 성능 향상으로 직결된다.

5. 수치적 안정성과 Segsum (Numerical Stability)

이론적으로 완벽한 SSD 알고리즘도 실제 구현 시에는 부동소수점 오차(Floating-point Error) 문제에 직면한다. 특히 1-반분리 행렬의 성분인 L_{ij} = \prod a_k는 0과 1 사이의 값들을 연속적으로 곱하는 연산이므로, 시퀀스가 길어질수록 값이 0으로 수렴하는 언더플로우(Underflow) 현상이 발생하기 쉽다.

5.1 로그 공간 연산 (Log-Space Computation)

이를 해결하기 위해 Mamba-2는 모든 누적 곱(Cumprod) 연산을 로그 공간(Log-space)에서의 누적 합(Cumsum)으로 변환하여 처리한다.
\prod_{k=j}^{i-1} a_k = \exp\left(\sum_{k=j}^{i-1} \log a_k\right)
이렇게 곱셈을 덧셈으로 변환하면 수치적 범위(Dynamic Range)가 넓어져 언더플로우를 방지할 수 있으며, GPU의 덧셈 연산 최적화를 활용할 수 있다.

5.2 Segsum 트릭과 치명적 상쇄 방지

하지만 단순히 로그 공간에서 cumsum을 구한 뒤 그 차이(cumsum[i] - cumsum[j])를 이용해 구간 합(Segment Sum)을 구하려 하면 “치명적 상쇄(Catastrophic Cancellation)” 문제가 발생할 수 있다. 매우 큰 두 수의 뺄셈 과정에서 유효 숫자가 소실되어 정확도가 급격히 떨어지는 현상이다.

Mamba-2는 이를 방지하기 위해 segsum이라는 특수화된 연산을 도입한다.7

  • Segsum: 단순히 전체에 대한 cumsum을 뺀 것이 아니라, 수치적으로 안정적인 방식으로 블록 내부에서 필요한 구간 합들을 독립적으로 계산한다.
  • 구현: segsum_unstable과 같은 나이브한 구현 대신, 안정성을 보장하는 커널을 사용하여 마스크 행렬 L을 생성한다. 이는 SSD 알고리즘이 이론적 성능을 실제 학습 결과의 정확도로 연결하는 데 있어 필수적인 저수준 최적화(Low-level Optimization)이다.

6. 이산화(Discretization)의 재해석과 유연성

기존의 구조적 SSM(S4, Mamba-1)은 연속 시간(Continuous-time) 모델을 정의한 후, 이를 이산화(Discretization)하는 과정(ZOH, Bilinear 변환 등)을 모델의 필수적인 부분으로 포함했다. 여기서 파라미터 \Delta (시간 간격)가 중요한 역할을 했다.

그러나 SSD 프레임워크에서는 이산화 단계가 더 이상 구조적인 필수 요소가 아니다. Mamba-2는 문제의 성격에 따라 두 가지 접근을 모두 허용한다.

  1. 직접 이산 학습: 모델을 h_t = A h_{t-1} + B x_t 형태의 이산 점화식 자체로 보고, 행렬 A와 B를 직접 최적화한다. 이는 텍스트와 같은 본질적인 이산 데이터(Tokenized Data)에 적합하다.
  2. 이산화 유지: 오디오나 시계열 센서 데이터와 같이 가변적인 샘플링 속도를 가지거나 연속적인 신호 처리가 필요한 경우, 기존처럼 연속 파라미터를 유지하고 이산화 단계를 SSD 알고리즘 앞단에 배치하여 유연성을 확보할 수 있다.7

결론적으로 5.2절의 행렬 믹서 이론과 블록 분해 알고리즘은 Mamba-2가 트랜스포머의 학습 효율성(병렬 처리)과 SSM의 추론 효율성(선형 시간)을 동시에 달성하게 만든 기술적 정점이다. 반분리 행렬이라는 수학적 도구를 통해 시퀀스 모델링을 재정의하고, 이를 블록 단위로 분해하여 최신 GPU 아키텍처에 최적화시킨 설계는, 단순히 기존 모델의 개선을 넘어 포스트 트랜스포머 시대의 새로운 연산 표준을 제시하고 있다.

7. 참고 자료

  1. State Space Duality (Mamba-2) Part II - The Theory | Tri Dao, https://tridao.me/blog/2024/mamba2-part2-theory/
  2. State Space Duality (Mamba-2) Part II - The Theory | Goomba Lab, https://goombalab.github.io/blog/2024/mamba2-part2-theory/
  3. Mamba-2: The ’Transform’ation of Mamba | by Utsavtiwari - Medium, https://medium.com/@utsavtiwari9936/mamba-2-the-transformation-of-mamba-125096294c51
  4. State Space Duality (Mamba-2) Part I - The Model | Goomba Lab, https://goombalab.github.io/blog/2024/mamba2-part1-model/
  5. Transformers are SSMs: Generalized Models and Efficient … - arXiv, https://arxiv.org/abs/2405.21060
  6. Mamba 2 | PDF | Matrix (Mathematics) | Tensor - Scribd, https://www.scribd.com/document/748683510/mamba2
  7. Mamba-2: Algorithms and Systems, https://pli.princeton.edu/blog/2024/mamba-2-algorithms-and-systems
  8. State Space Duality (Mamba-2) Part III - The Algorithm | Tri Dao, https://tridao.me/blog/2024/mamba2-part3-algorithm/
  9. State Space Duality (Mamba-2) Part IV - The Systems | Tri Dao, https://tridao.me/blog/2024/mamba2-part4-systems/